import smplx, torch
import os
import os.path as op
import numpy as np
import trimesh
from tqdm import tqdm
import pyvista as pv
from scipy.spatial.transform import Rotation as R

m_smplx = smplx.SMPLX(model_path=r'F:\workspaces\mmlab\cliff\data\body_models\smplx\models\smplx', gender='neutral', ext='npz', use_pca=False)
basdir = r'THuman2.1_Release\smplx'

file_idx = 532
rotate_i = 526

file_ls = sorted(os.listdir(basdir))[file_idx:file_idx+1]
for ids in tqdm(file_ls):
    params = np.load(op.join(basdir, ids, 'smplx_param.pkl'), allow_pickle=True)
    
    transl=torch.tensor(params['transl'], dtype=torch.float32),
    global_orient=torch.tensor(params['global_orient'], dtype=torch.float32),
    body_pose=torch.tensor(params['body_pose'], dtype=torch.float32),
    betas=torch.tensor(params['betas'], dtype=torch.float32),
    left_hand_pose=torch.tensor(params['left_hand_pose'], dtype=torch.float32),
    right_hand_pose=torch.tensor(params['right_hand_pose'], dtype=torch.float32),
    expression=torch.tensor(params['expression'], dtype=torch.float32),
    jaw_pose=torch.tensor(params['jaw_pose'], dtype=torch.float32),
    leye_pose=torch.tensor(params['leye_pose'], dtype=torch.float32),
    reye_pose=torch.tensor(params['reye_pose'], dtype=torch.float32)
    
    smpl_out = m_smplx.forward(
        transl=torch.tensor(params['transl'], dtype=torch.float32),
        global_orient=torch.tensor(params['global_orient'], dtype=torch.float32),
        body_pose=torch.tensor(params['body_pose'], dtype=torch.float32),
        betas=torch.tensor(params['betas'], dtype=torch.float32),
        left_hand_pose=torch.tensor(params['left_hand_pose'], dtype=torch.float32),
        right_hand_pose=torch.tensor(params['right_hand_pose'], dtype=torch.float32),
        expression=torch.tensor(params['expression'], dtype=torch.float32),
        jaw_pose=torch.tensor(params['jaw_pose'], dtype=torch.float32),
        leye_pose=torch.tensor(params['leye_pose'], dtype=torch.float32),
        reye_pose=torch.tensor(params['reye_pose'], dtype=torch.float32)
    )
    # vertices = smpl_out.vertices.detach().cpu().numpy()[0] * params['scale']
    # m_model = trimesh.Trimesh(vertices=vertices, faces=m_smplx.faces, process=False)
    # m_model.export('debug/%s.obj' % ids)
    vertices = smpl_out.vertices.detach().cpu().numpy()[0] * params['scale']
    
    ########
    # euler_angle = global_orient[0][0]
    # r = R.from_euler('xyz', euler_angle)
    euler_angle = [90, 0, 0]
    r = R.from_euler('xyz', euler_angle, degrees=True)
    
    rotation_matrix = r.as_matrix()
    rotation_matrix_inv = np.linalg.inv(rotation_matrix)
    
    if int(ids) < rotate_i:
        vertices_new = vertices
    else:
        vertices_new = np.matmul(rotation_matrix_inv, vertices.T).T
    
    # normalize
    min_x_vals = vertices_new[:, 0].min()
    max_x_vals = vertices_new[:, 0].max()
    min_y_vals = vertices_new[:, 1].min()
    max_y_vals = vertices_new[:, 1].max()
    min_z_vals = vertices_new[:, 2].min()
    max_z_vals = vertices_new[:, 2].max()
    # center = (max_vals + min_vals) / 2.0
    # scale = (max_vals - min_vals).max() / 2.0
    # vertices_new = (vertices_new - center) / scale
    print(f"\nmin: {min_x_vals}, {min_y_vals}, {min_z_vals}; max: {max_x_vals}, {max_y_vals}, {max_z_vals}")
    ########


    mesh = trimesh.Trimesh(vertices=vertices_new, faces=m_smplx.faces, process=False)
    # 创建PyVista PolyData对象
    vertices = mesh.vertices

    min_x_vals = vertices[:, 0].min()
    max_x_vals = vertices[:, 0].max()
    min_y_vals = vertices[:, 1].min()
    max_y_vals = vertices[:, 1].max()
    min_z_vals = vertices[:, 2].min()
    max_z_vals = vertices[:, 2].max()
    print(f"\nmin: {min_x_vals}, {min_y_vals}, {min_z_vals}; max: {max_x_vals}, {max_y_vals}, {max_z_vals}")

    faces = mesh.faces.reshape((-1, 3))
    faces = np.hstack([np.full((faces.shape[0], 1), 3), faces])  # PyVista faces格式要求
    polydata = pv.PolyData(vertices, faces)
    # 创建Plotter对象并绘制mesh
    plotter = pv.Plotter()
    color = 'lightblue'    #'lightblue'
    plotter.add_mesh(polydata, color=color, show_edges=False, edge_color='gray')
    # 添加坐标轴
    plotter.add_axes(interactive=True)
    # 设置初始相机位置以保持右手坐标系
    plotter.view_xy()  # 你可以根据需要调整为 view_xy, view_xz 或 view_yz

    plotter.show()

    exit()